#pragma once

#include "BitStream.h"
#include <cassert>
#include <algorithm>
#include <functional>

namespace phynet{
namespace net{

BitStream::BitStream()
	:m_currentBit(0)
	,m_bitSize(0)
{
	m_writingBuffer = &m_stackBuffer[0];
	m_readingBuffer = m_writingBuffer;
	m_bufferSize = sizeof(m_stackBuffer);
}

BitStream::BitStream(UInt32 i_size)
	:m_currentBit(0)
	,m_bitSize(0)
{
	m_internalBuffer.resize(i_size);
	m_writingBuffer = &m_internalBuffer.at(0);
	m_readingBuffer = m_writingBuffer;
	m_bufferSize = i_size;
}

BitStream::BitStream(const Byte* i_externalBuffer, UInt32 i_bufferSize)
:m_bitSize(i_bufferSize <<3)
,m_currentBit(0)
,m_readingBuffer(i_externalBuffer)
,m_writingBuffer(0)
,m_bufferSize(i_bufferSize)
{
}

void BitStream::Copy(const BitStream& i_stream)
{
	Clear();
	UInt32	copySize = min(i_stream.GetByteSize(), this->GetBufferSize());
	memcpy(m_writingBuffer,i_stream.m_readingBuffer,copySize);
	m_bitSize	 = copySize<<3;
	m_currentBit = min(i_stream.m_currentBit, m_bitSize);
}

void BitStream::Resize(const UInt32 i_size)
{
	m_internalBuffer.resize(i_size);
	m_writingBuffer = &m_internalBuffer.at(0);
	m_readingBuffer = m_writingBuffer;
	m_bufferSize = i_size;
	Clear();
}

void BitStream::WriteUInt8(UInt8 val)
{
	UInt8 rbits = GetCurrentByteOffset();
	UInt8 shift = (8 - rbits);
	UInt32 currentByte = GetCurrentByte();
	m_writingBuffer[currentByte] &= ~((1 << shift) -1);
	m_writingBuffer[currentByte] |= (val >> rbits);
	m_writingBuffer[++currentByte] = (val << shift);
	m_currentBit += 8;
	UpdateBitSize();
}

UInt8 BitStream::ReadUInt8() const
{
	UInt8 rbits = GetCurrentByteOffset();
	UInt8 shift = (8 - rbits);
	UInt8 val;
	UInt8 *v = (UInt8*)&val;
	UInt32 currentByte = GetCurrentByte();
	*v = (UInt8)(m_readingBuffer[currentByte] << rbits);
	*v |= (UInt8)(m_readingBuffer[++currentByte] >> shift);
	m_currentBit += 8;
	return val;
}

void BitStream::WriteUInt16(UInt16 val)
{
	UInt8 rbits = GetCurrentByteOffset();
	UInt8 shift = (8 - rbits);
	UInt32 currentByte = GetCurrentByte();
	UInt8 *v = (UInt8*)(&val);
	m_writingBuffer[currentByte] &= ~((1 << shift) -1);
	for(Int32 i=0; i < sizeof(UInt16); i++)
	{
		m_writingBuffer[currentByte]  |= (*v >> rbits);
		m_writingBuffer[++currentByte] = (*v << shift);
		++v;
	}
	m_currentBit += 16;
	UpdateBitSize();
}

UInt16 BitStream::ReadUInt16() const
{
	UInt8 rbit  = static_cast<UInt8>(m_currentBit & 0x7);
	UInt8 shift = (8 - rbit);
	UInt16 val;
	UInt8 *v = (UInt8 *)&val;
	UInt32 currentByte = GetCurrentByte();
	for(Int32 i=0; i < sizeof(short); i++)
	{
		*v = (UInt8)(m_readingBuffer[currentByte] << rbit);
		*v |= (UInt8)(m_readingBuffer[++currentByte] >> shift);
		v++;
	}
	m_currentBit += 16;
	return val;
}

void BitStream::WriteUInt32(UInt32 val)
{
	UInt8 rbits = GetCurrentByteOffset();
	UInt8 shift = (8 - rbits);
	UInt8 *v = (UInt8 *)&val;
	UInt32 currentByte = GetCurrentByte();
	m_writingBuffer[currentByte] &= ~((1 << shift) -1);
	for(Int32 i=0; i < sizeof(UInt32); i++)
	{
		m_writingBuffer[currentByte]  |= (*v >> rbits);
		m_writingBuffer[++currentByte] = (*v << shift);
		++v;
	}
	m_currentBit += 32;
	UpdateBitSize();
}

UInt32 BitStream::ReadUInt32() const
{
	UInt8 rbit  = static_cast<UInt8>(m_currentBit & 0x7);
	UInt8 shift = (8 - rbit);
	//8 bytes
	UInt32 val;
	UInt8 *v = (UInt8 *)&val;
	UInt32 currentByte = GetCurrentByte();
	for(Int32 i=0; i < sizeof(Int32); i++)
	{
		//first byte
		*v = (UInt8)(m_readingBuffer[currentByte] << rbit);
		*v |= (UInt8)(m_readingBuffer[++currentByte] >> shift);
		v++;
	}
	m_currentBit += 32;
	return val;
}

void BitStream::WriteUInt64(UInt64 val)
{
	UInt8 rbits = GetCurrentByteOffset();
	UInt8 shift = (8 - rbits);
	UInt8 *v = (UInt8 *)&val;
	UInt32 currentByte = GetCurrentByte();
	m_writingBuffer[currentByte] &= ~((1 << shift) -1);
	for(Int32 i=0; i < sizeof(UInt64); i++)
	{
		m_writingBuffer[currentByte]  |= (*v >> rbits);
		m_writingBuffer[++currentByte] = (*v << shift);
		++v;
	}
	m_currentBit += 64;
	UpdateBitSize();
}

UInt64 BitStream::ReadUInt64() const
{
	UInt8 rbit  = static_cast<UInt8>(m_currentBit & 0x7);
	UInt8 shift = (8 - rbit);
	UInt64	val;
	UInt8 *v = (UInt8 *)&val;
	UInt32 currentByte = GetCurrentByte();
	for(Int32 i=0; i < sizeof(UInt64); i++)
	{
		//first byte
		*v = (UInt8)(m_readingBuffer[currentByte] << rbit);
		*v |= (UInt8)(m_readingBuffer[++currentByte] >> shift);
		v++;
	}
	m_currentBit += 64;
	return val;
}

void BitStream::WriteReal(Real val)
{
	UInt8 rbits = GetCurrentByteOffset();
	UInt8 shift = (8 - rbits);
	UInt8 *v = (UInt8*)&val;
	UInt32 currentByte = GetCurrentByte();
	m_writingBuffer[currentByte] &= ~((1 << shift) -1);
	for(Int32 i=0; i < sizeof(Real); i++)
	{
		m_writingBuffer[currentByte]  |= (*v >> rbits);
		m_writingBuffer[++currentByte] = (*v << shift);
		++v;
	}
	m_currentBit += sizeof(Real)*8;
	UpdateBitSize();
}

Real BitStream::ReadReal() const
{
	UInt8 rbit  = static_cast<UInt8>(m_currentBit & 0x7);
	UInt8 shift = (8 - rbit);
	Real val;
	UInt8 *v = (UInt8 *)&val;
	UInt32 currentByte = GetCurrentByte();
	for(Int32 i=0; i < sizeof(Real); i++)
	{
		//first byte
		*v = (UInt8)(m_readingBuffer[currentByte] << rbit);
		*v |= (UInt8)(m_readingBuffer[++currentByte] >> shift);
		v++;
	}
	m_currentBit += sizeof(Real)*8;
	return val;
}

void BitStream::WriteBytes(const UInt8* data, Int32 size)
{
	assert(GetCurrentByte()+size<=GetBufferSize());
	assert(data || (size == 0));

	while(size>0)
	{
		WriteUInt8(*data);
		--size;
		++data;
	}
}

void BitStream::ReadBytes(UInt8* data, Int32 size) const
{
	assert(GetCurrentByte()+size<=GetBufferSize());
	assert(data);

	while(size>0)
	{
		*data = ReadUInt8();
		++data;
		--size;
	}
}

void BitStream::WriteBit(Bit the_bit)
{
	UInt32 currentByte = GetCurrentByte();
	UInt8 rbits = static_cast<UInt8>(m_currentBit & 0x7);
	m_writingBuffer[currentByte] &= (0xff << (8 - rbits));
	if(the_bit)
	{	
		m_writingBuffer[currentByte] |= static_cast<UInt8>(1 << (7 - rbits));
	}
	m_currentBit++;
	UpdateBitSize();
}

Bit  BitStream::ReadBit() const
{
	UInt32 currentByte = GetCurrentByte();
	UInt8 rbits = static_cast<UInt8>(m_currentBit & 0x7);
	UInt8 shift = 7 - rbits;
	Bit return_value = static_cast<Bit>((m_readingBuffer[currentByte] >> shift) & 0x1);
	m_currentBit++;
	return return_value;
}

void BitStream::WriteBits(const UInt8* data, UInt32 bitSize)
{
	assert(data);

	while (bitSize >= 8)
	{
		WriteUInt8(*data);
		++data;
		bitSize -=8 ;
	}

	if(bitSize!=0)
	{
		UInt8* out = m_writingBuffer + (m_currentBit >> 3);
		UInt8 currentWritingByte = *data;
		const UInt8 rbits = static_cast<UInt8>(m_currentBit & 0x7);
		const UInt8 shift = 8 - rbits;
		const UInt8 rshift = static_cast<UInt8>(8 - bitSize);

		if(bitSize <= shift)
		{
			if(rbits==0)
			{
				*out = 0;
			}
			currentWritingByte >>= rshift;
			currentWritingByte <<= shift - bitSize; 
			*out |= currentWritingByte;
		}
		else
		{
			*out |= currentWritingByte >> rbits;
			++out;
			currentWritingByte >>= rshift;
			currentWritingByte <<= rshift + shift;
			*out = currentWritingByte;
		}
		m_currentBit +=bitSize;
		UpdateBitSize();
	}
}

void BitStream::ReadBits(UInt8* data, UInt32 bitSize) const
{
	assert(data);

	while (bitSize>=8)
	{
		*data = ReadUInt8();
		++data;
		bitSize-=8;
	}

	assert(bitSize<8);

	if(bitSize!=0)
	{
		const UInt8* input = m_readingBuffer + (m_currentBit >> 3);
		const UInt8 rbits = static_cast<UInt8>(m_currentBit & 0x7);
		const UInt8 shift = 8 - rbits;

		*data = *input;
		if(bitSize <= shift)
		{
			*data >>= (shift - bitSize);
			*data <<= 8 - bitSize;
		}
		else
		{
			*data <<= rbits;
			++input;
			*data |= (*input) >> shift;
		}
		m_currentBit += bitSize;
	}
}

}} //namespaces